package com.methodreturnvaluehandler;

import com.annotation.FileName;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.ByteArrayHttpMessageConverter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.DelegatingServerHttpResponse;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.web.context.request.NativeWebRequest;
import org.springframework.web.method.support.HandlerMethodReturnValueHandler;
import org.springframework.web.method.support.ModelAndViewContainer;
import org.springframework.web.servlet.mvc.method.annotation.HttpEntityMethodProcessor;

import javax.servlet.http.HttpServletResponse;
import java.io.InputStream;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;

/**
 * 用于将controller返回的{@link InputStream}转换为文件下载
 * 如果一个controller返回 {@link byte[]} 那么会将其转换为 文件下载
 * 如果controller 同时标注了 {@link FileName} 注解，那么会使用 其中的值作为下载文件的名称
 * 参考
 *
 * @see HttpEntityMethodProcessor#handleReturnValue(java.lang.Object, org.springframework.core.MethodParameter, org.springframework.web.method.support.ModelAndViewContainer, org.springframework.web.context.request.NativeWebRequest)
 */
public class ByteArrayReturnValueHandler implements HandlerMethodReturnValueHandler {

    HttpMessageConverter<byte[]> byteArrayHttpMessageConverter = new ByteArrayHttpMessageConverter();

    @Override
    public boolean supportsReturnType(MethodParameter returnType) {
        return byte[].class.isAssignableFrom(returnType.getParameterType());
    }

    @Override
    public void handleReturnValue(Object returnValue, MethodParameter returnType, ModelAndViewContainer mavContainer, NativeWebRequest webRequest) throws Exception {

        if (returnValue == null) {
            return;
        }

        byte[] data = (byte[]) returnValue;

        HttpServletResponse response = webRequest.getNativeResponse(HttpServletResponse.class);

        response.setContentType(MediaType.APPLICATION_OCTET_STREAM_VALUE);

        response.setCharacterEncoding(StandardCharsets.UTF_8.name());

        FileName methodAnnotation = returnType.getMethodAnnotation(FileName.class);

        String fileName = "下载的文件";

        if (methodAnnotation != null) {
            fileName = methodAnnotation.value();
        }

        // 这里URLEncoder.encode可以防止中文乱码
        String encodedFileName = URLEncoder.encode(fileName, StandardCharsets.UTF_8.name()).replaceAll("\\+", "%20");
        response.setHeader("Content-Disposition", "attachment;filename=" + encodedFileName + ".txt");


        ServletServerHttpResponse httpResponse = new ServletServerHttpResponse(response);

        byteArrayHttpMessageConverter.write(data, MediaType.APPLICATION_OCTET_STREAM, httpResponse);

//      标记此请求已被处理，避免DispatcherServlet进行视图解析
        mavContainer.setRequestHandled(true);


    }
}
